//////////////////////////////////////////////////////////////////
// (c) Copyright 1998-2002 by Jeongnim Kim
//////////////////////////////////////////////////////////////////
//////////////////////////////////////////////////////////////////
//   Jeongnim Kim
//   National Center for Supercomputing Applications &
//   Materials Computation Center
//   University of Illinois, Urbana-Champaign
//   Urbana, IL 61801
//   e-mail: jnkim@ncsa.uiuc.edu
//   Tel:    217-244-6319 (NCSA) 217-333-3324 (MCC)
//
// Supported by 
//   National Center for Supercomputing Applications, UIUC
//   Materials Computation Center, UIUC
//   Department of Physics, Ohio State University
//   Ohio Supercomputer Center
//////////////////////////////////////////////////////////////////
// -*- C++ -*-
#ifndef OHMMS_CG2IS_TEMPLATE_H
#define OHMMS_CG2IS_TEMPLATE_H
#include "Utilities/Clock.h"

/*!\class CG_2I_S 
 * \brief CG minimization to minimize energy function
 * 2I-S overlap matrices for orthonormal basis functions
 * Entire bands are minimized.
 *
 * Instead of expression templates, decide to use functions with
 * explicit names. For dense matrix operations, blas 3 is used.
 *
 * operations to be provided:
 *   generic Y = Y + a*X                  axpy(a,X,Y) 
 *   generic C = x A*B + yC               xAB_plus_yC(x,A,B,y,C)
 *   generic C = x A*B^t + yC             xABt_plus_yC(x,A,B,y,C)
 *   generic C = x A*A^t + yC             xAAt_plus_yC(x,A,y,C)
 *   generic C = x (A*B^t + B*A^t)+yC     xABt_xBAt_plus_yC(x,A,B,y,C)
 *
 * virtual member function apply(x,y)
 * y <- H x for EigenState_t x and y
 */

template<class EigenStates_t, class DensityMatrix_t>
struct CG_2I_S {

  //typedef typename EigenStates_t::value_type::value_type value_type;
  typedef double value_type;
  int cgstep;
  int cgstep_min;
  int cgstep_max;
  int debug_level;
  int cgstep_acc;

  value_type cpuaccum[10];
  value_type etol, gtol;
  value_type gdotg, gdotg0, gdoth;

  EigenStates_t temp;
  EigenStates_t gr, cg, gr0;    //gradient, conjugate gradient
  DensityMatrix_t H, Q;           //overlap matrices
  DensityMatrix_t H1, H2, S1, S2; //overlap matrices

  CG_2I_S() { 
    cgstep = -1;
    cgstep_min = 2;
    cgstep_max = 100;
    etol = 1e-10;
    gtol = 1e-6;
    debug_level = 0;
    for(int i=0; i<10; i++) cpuaccum[i] = 0.0;
  }

  //clean up the vectors
  ~CG_2I_S() { }

  //any method needs to be inherited from this
  virtual 
  void apply(const EigenStates_t& a, EigenStates_t& b) = 0;

  //any method needs to be inherited from this
  virtual  
  void init(const EigenStates_t& psi) = 0;

  /*!\fn void densitymatrix
   *\brief perform A^t*Q A to caluclate a density matrix
   */
  EigenStates_t&
  Conj(const EigenStates_t& psi) {
    xAB_plus_yC(1.0, Q, psi, 0.0, temp);
    return temp;
  }

  value_type 
  minimize(EigenStates_t& psi, double shift= 0) {

    Pooma::Clock cgtimer;

    init(psi); // check the states
    value_type e0, eold, mu;
    mu = shift;
    bool notconverged = true;


    do {
      cgtimer.start();
      xAAt_plus_yC(-1.0,psi,0.0,Q); 
      //Q(i,j) =  psi(i,k)*psi(j,k) <=> Q = Psi*Transpose(Psi)
      addDiagonal(Q,2.0);
      cgtimer.stop();
      cpuaccum[0] += cgtimer.cpu_time();

      cgtimer.start();
      apply(psi,temp); // Psi_1 = \hat H |Psi>, for any \hat H to be minimized
      cgtimer.stop();
      cpuaccum[1] += cgtimer.cpu_time();

      if(fabs(mu) > etol) axpy(-mu,psi,temp); // generic temp = temp - mu*psi

      cgtimer.start();

      // H = <Psi|(h|Psi>) = <Psi|Psi_1>
      //H(i,j) = psi(i,k)*(\hat H psi(j,k)) <=> H = Psi*Transpose(\hat H Psi)
      xABt_plus_yC(1.0, psi, temp, 0.0, H); 

      // G = <Psi|h|Psi> |Psi> - (2I-<Psi|Psi>) h |Psi>
      xAB_plus_yC( 1.0, H, psi, 0.0, gr); // gr  = H|psi>
      xAB_plus_yC(-1.0, Q, temp, 1.0, gr);// gr -= Q|psi>
      cgtimer.stop();
      cpuaccum[2] += cgtimer.cpu_time();
    
      cgtimer.start();

      e0 = Dot(H,Q);      // Trace(H*transpose(Q)) for symmetric H and Q
      gdotg = Dot(gr,gr); // Trace <=> Dot
      cgtimer.stop();
      cpuaccum[3] += cgtimer.cpu_time();

      if(cgstep == 0) {
	copy(gr,cg);
      } else {
	//Polak-Ribiere (defualt)
	//gdoth = Dot(gr,gr0);
	//value_type gamma = (gdotg - gdoth)/gdotg0; 

	// so far, best with Tr((2I-S)H) functional
	//Fletcher-Reeve
	value_type gamma = gdotg/gdotg0;
	axpy(gamma,gr,cg); // generic y = y + ax

	//does not seem to work
	//gdotg0 = Dot(gr,gr0);
	//gdoth = Dot(gr,cg)-Dot(gr0,cg);
	//value_type gamma = -(gdotg-gdotg0)/(gdoth+1e-20);
	//axpy(gamma,gr,cg); // generic y = y + ax
	//gr0 = gr;
      }

      if(cgstep > cgstep_min) {
	if(fabs((eold-e0)/e0) < etol)  {
std::cout << " energy tolerance has been reached " << std::endl; 
	  notconverged = false;
        }
	if(gdotg < gtol) {
std::cout << " gradient tolerance has been reached " << std::endl; 
	  notconverged = false;
        }
      }

      gdotg0 = gdotg;// save <G|G>
      eold = e0;     // save E
      
      apply(cg,temp); // temp = \hat cg
      if(fabs(mu) > etol) axpy(-mu,cg,temp); // temp = temp -mu*cg

      xABt_plus_yC(1.0, temp, cg,0.0,H2);
      xABt_xBAt_plus_yC(1.0, psi, temp, 0.0, H1);

      xAAt_plus_yC(1.0, cg, 0.0, S2);
      xABt_xBAt_plus_yC(1.0, psi, cg, 0.0, S1);

      //rely on the fact that H's S's are symmetrized
      //value_type e1 = Trace(H1,Q)-Trace(H,S1);
      //value_type e2 = Trace(H2,Q)-Trace(H1,S1)-Trace(H,S2);
      value_type e1 = Dot(H1,Q)-Dot(H,S1);
      value_type e2 = Dot(H2,Q)-Dot(H1,S1)-Dot(H,S2);

      // E[lamda] = E_0 + lambda*e1 + lambda*lambda*e2
      // dE/dlambda = 0 -> lambda = -e1/(2*e2)
      value_type lambda = -e1/e2/2.0;
      value_type delE = lambda*(e1+lambda*e2);
      if(debug_level > 1) {
	cout << "ene = " << e0 
	     << "  ene_new [lambda =" << lambda << "]= " << delE << endl;
      }
      if(lambda < 0 || lambda > 0.5) {
	cg = gr;//equivalent to the restarting cg
	lambda = 1e-4;
      }
      cgstep++;
      axpy(lambda,cg,psi);
    } while(cgstep < cgstep_max && notconverged);

    xAAt_plus_yC(-1.0,psi,0.0,Q);
    addDiagonal(Q,2.0);

    // Psi_1 = h |Psi>
    apply(psi,temp);

    // H = <Psi|(h|Psi>) = <Psi|Psi_1>
    xABt_plus_yC(1.0, psi, temp, 0.0, H); 
    e0 = Dot(H,Q); // again, H and Q are symmetric matries
    //e0 = Trace(H,Q);

    cgstep_acc+=cgstep;
    if(debug_level > 0) {
      cout << "ene = " << e0 << " cgstep = " << cgstep_acc  << " "
	   << cpuaccum[0] << " " << cpuaccum[1] << " " << cpuaccum[2] << " " 
	   << cpuaccum[3] << endl << endl; 
    }
    cgstep = 0; // reset cgstep to start anew
    return e0;
  }

};

#endif
/***************************************************************************
 * $RCSfile: CG_2I_S_Solver.h,v $   $Author: jnkim $
 * $Revision: 1.1.1.1 $   $Date: 2004/07/23 21:04:02 $
 * $Id: CG_2I_S_Solver.h,v 1.1.1.1 2004/07/23 21:04:02 jnkim Exp $ 
 ***************************************************************************/
